{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.random import default_rng\n",
    "import random\n",
    "import collections\n",
    "import re\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_length_encoder = 512\n",
    "max_seq_length_decoder = 128\n",
    "masked_lm_prob = 0.2\n",
    "max_predictions_per_seq = int(masked_lm_prob * max_seq_length_encoder)\n",
    "do_whole_word_mask = True\n",
    "EOS_ID = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "MaskedLmInstance = collections.namedtuple(\n",
    "    'MaskedLmInstance', ['index', 'label']\n",
    ")\n",
    "\n",
    "\n",
    "class TrainingInstance(object):\n",
    "    \"\"\"A single training instance (sentence pair).\"\"\"\n",
    "\n",
    "    def __init__(self, tokens, tokens_y, masked_lm_positions, masked_lm_labels):\n",
    "        self.tokens = tokens\n",
    "        self.tokens_y = tokens_y\n",
    "        self.masked_lm_positions = masked_lm_positions\n",
    "        self.masked_lm_labels = masked_lm_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sliding(strings, n = 5):\n",
    "    results = []\n",
    "    for i in range(len(strings) - n):\n",
    "        results.append(strings[i : i + n])\n",
    "    return results\n",
    "\n",
    "\n",
    "def _get_ngrams(n, text):\n",
    "    ngram_set = set()\n",
    "    text_length = len(text)\n",
    "    max_index_ngram_start = text_length - n\n",
    "    for i in range(max_index_ngram_start + 1):\n",
    "        ngram_set.add(tuple(text[i : i + n]))\n",
    "    return ngram_set\n",
    "\n",
    "\n",
    "def _get_word_ngrams(n, sentences):\n",
    "    assert len(sentences) > 0\n",
    "    assert n > 0\n",
    "\n",
    "    words = sum(sentences, [])\n",
    "    return _get_ngrams(n, words)\n",
    "\n",
    "\n",
    "def cal_rouge(evaluated_ngrams, reference_ngrams):\n",
    "    reference_count = len(reference_ngrams)\n",
    "    evaluated_count = len(evaluated_ngrams)\n",
    "\n",
    "    overlapping_ngrams = evaluated_ngrams.intersection(reference_ngrams)\n",
    "    overlapping_count = len(overlapping_ngrams)\n",
    "\n",
    "    if evaluated_count == 0:\n",
    "        precision = 0.0\n",
    "    else:\n",
    "        precision = overlapping_count / evaluated_count\n",
    "\n",
    "    if reference_count == 0:\n",
    "        recall = 0.0\n",
    "    else:\n",
    "        recall = overlapping_count / reference_count\n",
    "\n",
    "    f1_score = 2.0 * ((precision * recall) / (precision + recall + 1e-8))\n",
    "    return {'f': f1_score, 'p': precision, 'r': recall}\n",
    "\n",
    "\n",
    "def _rouge_clean(s):\n",
    "    return re.sub(r'[^a-zA-Z0-9 ]', '', s)\n",
    "\n",
    "\n",
    "def get_rouges(strings, n = 1):\n",
    "    rouges = []\n",
    "    for i in range(len(strings)):\n",
    "        abstract = strings[i]\n",
    "        doc_sent_list = [strings[k] for k in range(len(strings)) if k != i]\n",
    "        sents = _rouge_clean(' '.join(doc_sent_list)).split()\n",
    "        abstract = _rouge_clean(abstract).split()\n",
    "        evaluated_1grams = _get_word_ngrams(n, [sents])\n",
    "        reference_1grams = _get_word_ngrams(n, [abstract])\n",
    "        rouges.append(cal_rouge(evaluated_1grams, reference_1grams)['f'])\n",
    "    return rouges\n",
    "\n",
    "\n",
    "# Principal Select top-m scored sentences according to importance.\n",
    "# As a proxy for importance we compute ROUGE1-F1 (Lin, 2004) between the sentence and the rest of the document\n",
    "def get_rouge(strings, top_k = 1, minlen = 4):\n",
    "    rouges = get_rouges(strings)\n",
    "    s = np.argsort(rouges)[::-1]\n",
    "    s = [i for i in s if len(strings[i].split()) >= minlen]\n",
    "    return s[:top_k]\n",
    "\n",
    "\n",
    "# Random Uniformly select m sentences at random.\n",
    "def get_random(strings, top_k = 1):\n",
    "    return rng.choice(len(strings), size = top_k, replace = False)\n",
    "\n",
    "\n",
    "# Lead Select the first m sentences.\n",
    "def get_lead(strings, top_k = 1):\n",
    "    return [i for i in range(top_k)]\n",
    "\n",
    "\n",
    "def combine(l):\n",
    "    r = []\n",
    "    for s in l:\n",
    "        if s[-1] != '.':\n",
    "            if s in ['[MASK]', '[MASK2]']:\n",
    "                e = ' .'\n",
    "            else:\n",
    "                e = '.'\n",
    "            s = s + e\n",
    "        r.append(s)\n",
    "    return ' '.join(r)\n",
    "\n",
    "\n",
    "def is_number_regex(s):\n",
    "    if re.match('^\\d+?\\.\\d+?$', s) is None:\n",
    "        return s.isdigit()\n",
    "    return True\n",
    "\n",
    "\n",
    "def reject(token):\n",
    "    t = token.replace('##', '')\n",
    "    if is_number_regex(t):\n",
    "        return True\n",
    "    if t.startswith('RM'):\n",
    "        return True\n",
    "    if token in '!{<>}:;.,\"\\'':\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def create_masked_lm_predictions(\n",
    "    tokens,\n",
    "    vocab_words,\n",
    "    rng,\n",
    "):\n",
    "    \"\"\"Creates the predictions for the masked LM objective.\"\"\"\n",
    "\n",
    "    cand_indexes = []\n",
    "    for (i, token) in enumerate(tokens):\n",
    "        if token == '[CLS]' or token == '[SEP]' or token == '[MASK2]':\n",
    "            continue\n",
    "        if reject(token):\n",
    "            continue\n",
    "        if (\n",
    "            do_whole_word_mask\n",
    "            and len(cand_indexes) >= 1\n",
    "            and token.startswith('##')\n",
    "        ):\n",
    "            cand_indexes[-1].append(i)\n",
    "        else:\n",
    "            cand_indexes.append([i])\n",
    "\n",
    "    rng.shuffle(cand_indexes)\n",
    "\n",
    "    output_tokens = list(tokens)\n",
    "\n",
    "    num_to_predict = min(\n",
    "        max_predictions_per_seq,\n",
    "        max(1, int(round(len(tokens) * masked_lm_prob))),\n",
    "    )\n",
    "\n",
    "    masked_lms = []\n",
    "    covered_indexes = set()\n",
    "    for index_set in cand_indexes:\n",
    "        if len(masked_lms) >= num_to_predict:\n",
    "            break\n",
    "        if len(masked_lms) + len(index_set) > num_to_predict:\n",
    "            continue\n",
    "        is_any_index_covered = False\n",
    "        for index in index_set:\n",
    "            if index in covered_indexes:\n",
    "                is_any_index_covered = True\n",
    "                break\n",
    "        if is_any_index_covered:\n",
    "            continue\n",
    "        for index in index_set:\n",
    "            covered_indexes.add(index)\n",
    "\n",
    "            masked_token = None\n",
    "            # 80% of the time, replace with [MASK]\n",
    "            if rng.random() < 0.8:\n",
    "                masked_token = '[MASK]'\n",
    "            else:\n",
    "                # 10% of the time, keep original\n",
    "                if rng.random() < 0.5:\n",
    "                    masked_token = tokens[index]\n",
    "                # 10% of the time, replace with random word\n",
    "                else:\n",
    "                    masked_token = vocab_words[\n",
    "                        np.random.randint(0, len(vocab_words) - 1)\n",
    "                    ]\n",
    "\n",
    "            output_tokens[index] = masked_token\n",
    "\n",
    "            masked_lms.append(\n",
    "                MaskedLmInstance(index = index, label = tokens[index])\n",
    "            )\n",
    "    assert len(masked_lms) <= num_to_predict\n",
    "    masked_lms = sorted(masked_lms, key = lambda x: x.index)\n",
    "\n",
    "    masked_lm_positions = []\n",
    "    masked_lm_labels = []\n",
    "    for p in masked_lms:\n",
    "        masked_lm_positions.append(p.index)\n",
    "        masked_lm_labels.append(p.label)\n",
    "\n",
    "    return (output_tokens, masked_lm_positions, masked_lm_labels)\n",
    "\n",
    "\n",
    "def get_feature(x, y, tokenizer, vocab_words, rng, dedup_factor = 5, **kwargs):\n",
    "    tokens = tokenizer.tokenize(x)\n",
    "    if len(tokens) > (max_seq_length_encoder - 2):\n",
    "        tokens = tokens[:max_seq_length_encoder - 2]\n",
    "    \n",
    "    if '[MASK2]' not in tokens:\n",
    "        return []\n",
    "\n",
    "    tokens = ['[CLS]'] + tokens + ['[SEP]']\n",
    "    \n",
    "    tokens_y = tokenizer.tokenize(y)\n",
    "    if len(tokens_y) > (max_seq_length_decoder - 1):\n",
    "        tokens_y = tokens_y[:max_seq_length_decoder - 1]\n",
    "        \n",
    "    tokens_y = tokenizer.convert_tokens_to_ids(tokens_y)\n",
    "    tokens_y = tokens_y + [EOS_ID]\n",
    "    results = []\n",
    "    for i in range(dedup_factor):\n",
    "        output_tokens, masked_lm_positions, masked_lm_labels = create_masked_lm_predictions(\n",
    "            tokens, vocab_words, rng, **kwargs\n",
    "        )\n",
    "        output_tokens = tokenizer.convert_tokens_to_ids(output_tokens)\n",
    "        masked_lm_labels = tokenizer.convert_tokens_to_ids(masked_lm_labels)\n",
    "        t = TrainingInstance(\n",
    "            output_tokens, tokens_y, masked_lm_positions, masked_lm_labels\n",
    "        )\n",
    "        results.append(t)\n",
    "    return results\n",
    "\n",
    "\n",
    "def group_doc(data):\n",
    "    results, result = [], []\n",
    "    for i in data:\n",
    "        if not len(i) and len(result):\n",
    "            results.append(result)\n",
    "            result = []\n",
    "        else:\n",
    "            result.append(i)\n",
    "\n",
    "    if len(result):\n",
    "        results.append(result)\n",
    "    return results\n",
    "\n",
    "\n",
    "def create_int_feature(values):\n",
    "    feature = tf.train.Feature(\n",
    "        int64_list = tf.train.Int64List(value = list(values))\n",
    "    )\n",
    "    return feature\n",
    "\n",
    "\n",
    "def create_float_feature(values):\n",
    "    feature = tf.train.Feature(\n",
    "        float_list = tf.train.FloatList(value = list(values))\n",
    "    )\n",
    "    return feature\n",
    "\n",
    "\n",
    "def write_instance_to_example_file(\n",
    "    instances,\n",
    "    output_file\n",
    "):\n",
    "    writer = tf.python_io.TFRecordWriter(output_file)\n",
    "    for (inst_index, instance) in enumerate(instances):\n",
    "        input_ids = list(instance.tokens)\n",
    "        target_ids = list(instance.tokens_y)\n",
    "        while len(input_ids) < max_seq_length_encoder:\n",
    "            input_ids.append(0)\n",
    "            target_ids.append(0)\n",
    "        masked_lm_positions = list(instance.masked_lm_positions)\n",
    "        masked_lm_ids = list(instance.masked_lm_labels)\n",
    "        masked_lm_weights = [1.0] * len(masked_lm_ids)\n",
    "        while len(masked_lm_positions) < max_predictions_per_seq:\n",
    "            masked_lm_positions.append(0)\n",
    "            masked_lm_ids.append(0)\n",
    "            masked_lm_weights.append(0.0)\n",
    "\n",
    "        features = collections.OrderedDict()\n",
    "        features['input_ids'] = create_int_feature(input_ids)\n",
    "        features['targets_ids'] = create_int_feature(target_ids)\n",
    "        features['masked_lm_positions'] = create_int_feature(\n",
    "            masked_lm_positions\n",
    "        )\n",
    "        features['masked_lm_ids'] = create_int_feature(masked_lm_ids)\n",
    "        features['masked_lm_weights'] = create_float_feature(masked_lm_weights)\n",
    "        tf_example = tf.train.Example(\n",
    "            features = tf.train.Features(feature = features)\n",
    "        )\n",
    "        writer.write(tf_example.SerializeToString())\n",
    "\n",
    "    tf.logging.info('Wrote %d total instances', inst_index)\n",
    "\n",
    "\n",
    "def process_documents(\n",
    "    file,\n",
    "    output_file,\n",
    "    tokenizer,\n",
    "    min_slide = 5,\n",
    "    max_slide = 13,\n",
    "    dedup_mask = 2,\n",
    "):\n",
    "    with open(file) as fopen:\n",
    "        data = fopen.read().split('\\n')\n",
    "    rng = default_rng()\n",
    "    vocab_words = list(tokenizer.vocab.keys())\n",
    "    grouped = group_doc(data)\n",
    "    results = []\n",
    "    for s in range(min_slide, max_slide, 1):\n",
    "        for r in tqdm(grouped):\n",
    "            slided = sliding(r, s)\n",
    "            X, Y = [], []\n",
    "            for i in range(len(slided)):\n",
    "                try:\n",
    "                    strings = slided[i]\n",
    "                    rouge_ = get_rouge(strings)\n",
    "                    y = strings[rouge_[0]]\n",
    "                    strings[rouge_[0]] = '[MASK2]'\n",
    "                    x = combine(strings)\n",
    "                    result = get_feature(\n",
    "                        x,\n",
    "                        y,\n",
    "                        tokenizer,\n",
    "                        vocab_words,\n",
    "                        rng,\n",
    "                        dedup_factor = dedup_mask,\n",
    "                    )\n",
    "                    results.extend(result)\n",
    "                except:\n",
    "                    pass\n",
    "\n",
    "    write_instance_to_example_file(results, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/pure-text/tokenization.py:135: The name tf.gfile.GFile is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import tokenization\n",
    "\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file = 'pegasus.wordpiece', do_lower_case=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = 'dumping-cleaned-news.txt'\n",
    "output_file = 'news.tfrecord'\n",
    "min_slide = 5\n",
    "max_slide = 13\n",
    "dedup_mask = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(file) as fopen:\n",
    "    data = fopen.read().split('\\n')\n",
    "rng = default_rng()\n",
    "vocab_words = list(tokenizer.vocab.keys())\n",
    "grouped = group_doc(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:06<00:00, 15.92it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 13.48it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 13.44it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 12.69it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 12.66it/s]\n",
      "100%|██████████| 100/100 [00:08<00:00, 12.23it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 13.17it/s]\n",
      "100%|██████████| 100/100 [00:07<00:00, 13.59it/s]\n"
     ]
    }
   ],
   "source": [
    "results = []\n",
    "for s in range(min_slide, max_slide, 1):\n",
    "    for r in tqdm(grouped[:100]):\n",
    "        slided = sliding(r, s)\n",
    "        X, Y = [], []\n",
    "        for i in range(len(slided)):\n",
    "            try:\n",
    "                strings = slided[i]\n",
    "                rouge_ = get_rouge(strings)\n",
    "                y = strings[rouge_[0]]\n",
    "                strings[rouge_[0]] = '[MASK2]'\n",
    "                x = combine(strings)\n",
    "                result = get_feature(\n",
    "                    x,\n",
    "                    y,\n",
    "                    tokenizer,\n",
    "                    vocab_words,\n",
    "                    rng,\n",
    "                    dedup_factor = dedup_mask,\n",
    "                )\n",
    "                results.extend(result)\n",
    "            except:\n",
    "                pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Wrote 34829 total instances\n"
     ]
    }
   ],
   "source": [
    "write_instance_to_example_file(results, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm {output_file}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
